#include "Cut.h"

int cut(const char *src, const char *dst,double startTime,double endTime)
{
    // 1.处理参数
    av_log_set_level(AV_LOG_DEBUG);
    int ret = -1;
    int index = -1;

    AVFormatContext* pFmtCtx = NULL;
    AVFormatContext* oFmtCtx = NULL;
    const AVOutputFormat* outFmt = NULL;

    AVPacket pkt;

    int *stream_map = NULL;
    int stream_idx = 0;

    int64_t* dts_start_time = NULL;
    int64_t* pts_start_time = NULL;


    // 2.打开多媒体文件
    ret = avformat_open_input(&pFmtCtx,src,NULL,NULL);

    if(ret < 0)
    {
        av_log(NULL,AV_LOG_ERROR,"%s\n",av_err2str(ret));
        return -1;
    }

    // 3.打开目的文件的上下文
    avformat_alloc_output_context2(&oFmtCtx,NULL,NULL,dst);
    if(!oFmtCtx)
    {
        av_log(NULL,AV_LOG_ERROR,"no memory\n");
        goto _ERROR;
    }
    stream_map = av_calloc(pFmtCtx->nb_streams,sizeof(int));
    if(!stream_map)
    {
        av_log(NULL,AV_LOG_ERROR,"no memory\n");
        goto _ERROR;
    }

    for(int i = 0;i < pFmtCtx->nb_streams;i++)
    {
        AVStream* outStream = NULL;
        AVStream* inStream = pFmtCtx->streams[i];
        AVCodecParameters* inCodecPar = inStream->codecpar;
        if( inCodecPar->codec_type != AVMEDIA_TYPE_AUDIO&&
            inCodecPar->codec_type != AVMEDIA_TYPE_VIDEO&&
            inCodecPar->codec_type != AVMEDIA_TYPE_SUBTITLE)
        {
            stream_map[i] = -1;
            continue;
        }

        stream_map[i] = stream_idx++;
        // 4.为目的文件，创建一个新的音频流
        outStream = avformat_new_stream(oFmtCtx,NULL);

        if(!outStream)
        {
            av_log(oFmtCtx,AV_LOG_ERROR,"no memory\n");
            goto _ERROR;
        }

        avcodec_parameters_copy(outStream->codecpar,inStream->codecpar);
        outStream->codecpar->codec_tag = 0;
    }

    // 绑定
    ret = avio_open2(&oFmtCtx->pb,dst,AVIO_FLAG_WRITE,NULL,NULL);
    if(ret < 0)
    {
        av_log(oFmtCtx,AV_LOG_ERROR,"%s\n",av_err2str(ret));
        goto _ERROR;
    }
    // 7.写多媒体文件头到目标文件中
    ret = avformat_write_header(oFmtCtx,NULL);

    if(ret < 0 )
    {
        av_log(oFmtCtx,AV_LOG_ERROR,"%s\n",av_err2str(ret));
        goto _ERROR;
    }

    // 视频不能严格按照指定的时间截取？为什么？i帧？
    ret = av_seek_frame(pFmtCtx,-1,startTime * AV_TIME_BASE,AVSEEK_FLAG_BACKWARD);
    if(ret <0)
    {
        av_log(oFmtCtx,AV_LOG_ERROR,"%s\n",av_err2str(ret));
        goto _ERROR;
    }

    dts_start_time = av_calloc(pFmtCtx->nb_streams,sizeof(int64_t));
    for(int t =0;t < pFmtCtx->nb_streams;t++)
    {
        dts_start_time[t] = -1;
    }
    pts_start_time = av_calloc(pFmtCtx->nb_streams,sizeof(int64_t));
    for(int t =0;t < pFmtCtx->nb_streams;t++)
    {
        pts_start_time[t] = -1;
    }

    // 8.从源多媒体文件中读到视频/音频/字幕数据目的文件中
    while(av_read_frame(pFmtCtx,&pkt) >= 0)
    {
        AVStream* inStream,*outStream;

        if(dts_start_time[pkt.stream_index] == -1 && pkt.dts > 0)
        {
            dts_start_time[pkt.stream_index] = pkt.dts;
        }

        if(pts_start_time[pkt.stream_index] == -1 && pkt.pts > 0)
        {
            pts_start_time[pkt.stream_index] = pkt.pts;
        }

        inStream = pFmtCtx->streams[pkt.stream_index];

        if(av_q2d(inStream->time_base )* pkt.pts > endTime)
        {
            av_log(oFmtCtx,AV_LOG_INFO,"success!\n");
            break;
        }

        // if(pkt.stream_index < 0) // 有问题！！！为什么？
        if(stream_map[pkt.stream_index] < 0)
        {
            av_packet_unref(&pkt);
            continue;
        }
        pkt.pts = pkt.pts - pts_start_time[pkt.stream_index];
        pkt.dts = pkt.dts - dts_start_time[pkt.stream_index];
        if(pkt.dts > pkt.pts)
        {
            pkt.pts = pkt.dts;
        }

        pkt.stream_index = stream_map[pkt.stream_index];

        outStream = oFmtCtx->streams[pkt.stream_index];

        av_packet_rescale_ts(&pkt,inStream->time_base,outStream->time_base);

        // 起始时间

        pkt.pos = -1;
        av_interleaved_write_frame(oFmtCtx,&pkt);
        av_packet_unref(&pkt);
    }

    // 9.写多媒体文件尾到文件中
    av_write_trailer(oFmtCtx);

// 10.资源释放
_ERROR:
    if(pFmtCtx)
    {
        avformat_close_input(&pFmtCtx);
        pFmtCtx = NULL;
    }
    if(oFmtCtx && oFmtCtx->pb)
    {
        avio_close(oFmtCtx->pb);
    }
    if(oFmtCtx)
    {
        avformat_free_context(oFmtCtx);
        oFmtCtx = NULL;
    }
    if(stream_map)
    {
        av_free(stream_map);
    }

    if(dts_start_time)
    {
        av_free(dts_start_time);
    }

    if(pts_start_time)
    {
        av_free(pts_start_time);
    }

    return 0;
}

int cutDeepseek(const char *src, const char *dst,double startTime,double endTime)
{
    av_log_set_level(AV_LOG_DEBUG);
    int ret = -1;

    AVFormatContext* pFmtCtx = NULL;
    AVFormatContext* oFmtCtx = NULL;
    AVPacket pkt;

    int *stream_map = NULL;
    int stream_idx = 0;

    int64_t* start_dts = NULL;
    int64_t* start_pts = NULL;
    int *start_time_set = NULL;

    // 打开输入文件
    ret = avformat_open_input(&pFmtCtx, src, NULL, NULL);
    if(ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to open input: %s\n", av_err2str(ret));
        return -1;
    }

    // 获取流信息
    ret = avformat_find_stream_info(pFmtCtx, NULL);
    if(ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to find stream info: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 创建输出上下文
    avformat_alloc_output_context2(&oFmtCtx, NULL, NULL, dst);
    if(!oFmtCtx) {
        av_log(NULL, AV_LOG_ERROR, "Failed to create output context\n");
        goto _ERROR;
    }

    // 分配流映射数组
    stream_map = av_calloc(pFmtCtx->nb_streams, sizeof(int));
    start_dts = av_calloc(pFmtCtx->nb_streams, sizeof(int64_t));
    start_pts = av_calloc(pFmtCtx->nb_streams, sizeof(int64_t));
    start_time_set = av_calloc(pFmtCtx->nb_streams, sizeof(int));

    if(!stream_map || !start_dts || !start_pts || !start_time_set) {
        av_log(NULL, AV_LOG_ERROR, "Memory allocation failed\n");
        goto _ERROR;
    }

    // 初始化
    for(int i = 0; i < pFmtCtx->nb_streams; i++) {
        stream_map[i] = -1;
        start_dts[i] = -1;
        start_pts[i] = -1;
        start_time_set[i] = 0;
    }

    // 创建输出流
    for(int i = 0; i < pFmtCtx->nb_streams; i++) {
        AVStream* inStream = pFmtCtx->streams[i];
        AVCodecParameters* inCodecPar = inStream->codecpar;

        // 只处理音视频和字幕流
        if(inCodecPar->codec_type != AVMEDIA_TYPE_AUDIO &&
            inCodecPar->codec_type != AVMEDIA_TYPE_VIDEO &&
            inCodecPar->codec_type != AVMEDIA_TYPE_SUBTITLE) {
            continue;
        }

        AVStream* outStream = avformat_new_stream(oFmtCtx, NULL);
        if(!outStream) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to create output stream\n");
            goto _ERROR;
        }

        ret = avcodec_parameters_copy(outStream->codecpar, inCodecPar);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to copy codec parameters: %s\n", av_err2str(ret));
            goto _ERROR;
        }
        outStream->codecpar->codec_tag = 0;

        // 复制时间基
        outStream->time_base = inStream->time_base;

        stream_map[i] = stream_idx++;
    }

    // 打开输出文件
    if(!(oFmtCtx->oformat->flags & AVFMT_NOFILE)) {
        ret = avio_open(&oFmtCtx->pb, dst, AVIO_FLAG_WRITE);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to open output file: %s\n", av_err2str(ret));
            goto _ERROR;
        }
    }

    // 写入文件头
    ret = avformat_write_header(oFmtCtx, NULL);
    if(ret < 0) {
        av_log(oFmtCtx, AV_LOG_ERROR, "Failed to write header: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 定位到起始时间（使用精确搜索）
    int64_t seek_target = (int64_t)(startTime * AV_TIME_BASE);
    ret = avformat_seek_file(pFmtCtx, -1, INT64_MIN, seek_target, INT64_MAX, 0);
    if(ret < 0) {
        av_log(oFmtCtx, AV_LOG_WARNING, "Seek failed, trying alternative method: %s\n", av_err2str(ret));
        // 备用搜索方法
        ret = av_seek_frame(pFmtCtx, -1, seek_target, AVSEEK_FLAG_BACKWARD);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "All seek methods failed: %s\n", av_err2str(ret));
            goto _ERROR;
        }
    }

    // 读取并处理数据包
    while(av_read_frame(pFmtCtx, &pkt) >= 0) {
        int stream_index = pkt.stream_index;

        // 跳过未映射的流
        if(stream_map[stream_index] < 0) {
            av_packet_unref(&pkt);
            continue;
        }

        AVStream* inStream = pFmtCtx->streams[stream_index];
        double packet_pts = pkt.pts * av_q2d(inStream->time_base);

        // 检查是否超过结束时间
        if(packet_pts > endTime) {
            av_packet_unref(&pkt);
            break;
        }

        // 跳过开始时间之前的数据包
        if(packet_pts < startTime) {
            av_packet_unref(&pkt);
            continue;
        }

        // 记录起始时间戳
        if(!start_time_set[stream_index]) {
            start_dts[stream_index] = pkt.dts;
            start_pts[stream_index] = pkt.pts;
            start_time_set[stream_index] = 1;
        }

        // 调整时间戳
        if(start_dts[stream_index] != -1) {
            pkt.dts -= start_dts[stream_index];
        }
        if(start_pts[stream_index] != -1) {
            pkt.pts -= start_pts[stream_index];
        }

        // 确保PTS >= DTS（处理B帧）
        if(pkt.pts < pkt.dts) {
            pkt.pts = pkt.dts;
        }

        // 更新流索引
        pkt.stream_index = stream_map[stream_index];
        AVStream* outStream = oFmtCtx->streams[pkt.stream_index];

        // 时间戳缩放
        av_packet_rescale_ts(&pkt, inStream->time_base, outStream->time_base);

        pkt.pos = -1;
        pkt.duration = av_rescale_q(pkt.duration, inStream->time_base, outStream->time_base);

        // 写入数据包
        ret = av_interleaved_write_frame(oFmtCtx, &pkt);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to write frame: %s\n", av_err2str(ret));
            av_packet_unref(&pkt);
            goto _ERROR;
        }

        av_packet_unref(&pkt);
    }

    // 写入文件尾
    av_write_trailer(oFmtCtx);
    av_log(oFmtCtx, AV_LOG_INFO, "Cut successfully completed!\n");

_ERROR:
    // 资源清理
    if(pFmtCtx) {
        avformat_close_input(&pFmtCtx);
    }
    if(oFmtCtx) {
        if(!(oFmtCtx->oformat->flags & AVFMT_NOFILE) && oFmtCtx->pb) {
            avio_closep(&oFmtCtx->pb);
        }
        avformat_free_context(oFmtCtx);
    }
    av_freep(&stream_map);
    av_freep(&start_dts);
    av_freep(&start_pts);
    av_freep(&start_time_set);

    return (ret < 0) ? -1 : 0;
}

int reCutDeepseek(const char *src, const char *dst, double startTime, double endTime)
{
    // 1.处理参数
    av_log_set_level(AV_LOG_DEBUG);
    int ret = -1;

    AVFormatContext* pFmtCtx = NULL;
    AVFormatContext* oFmtCtx = NULL;

    AVPacket pkt;

    int *stream_map = NULL;
    int stream_idx = 0;

    int64_t global_start_pts = -1;
    int *stream_enabled = NULL;
    int video_stream_index = -1;
    int audio_stream_index = -1;

    // 2.打开多媒体文件
    ret = avformat_open_input(&pFmtCtx, src, NULL, NULL);
    if(ret < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "Failed to open input: %s\n", av_err2str(ret));
        return -1;
    }

    // 获取流信息（重要：缺少这一步）
    ret = avformat_find_stream_info(pFmtCtx, NULL);
    if(ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Failed to find stream info: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 根据文件扩展名确定输出格式
    const char* ext = strrchr(dst, '.');
    const char* format_name = NULL;

    if (ext) {
        if (strcmp(ext, ".mp4") == 0 || strcmp(ext, ".MP4") == 0) {
            format_name = "mp4";
        } else if (strcmp(ext, ".avi") == 0 || strcmp(ext, ".AVI") == 0) {
            format_name = "avi";
        } else if (strcmp(ext, ".mkv") == 0 || strcmp(ext, ".MKV") == 0) {
            format_name = "matroska";
        }
    }

    // 3.打开目的文件的上下文
    if (format_name) {
        ret = avformat_alloc_output_context2(&oFmtCtx, NULL, format_name, dst);
    } else {
        ret = avformat_alloc_output_context2(&oFmtCtx, NULL, NULL, dst);
    }

    if(!oFmtCtx || ret < 0)
    {
        av_log(NULL, AV_LOG_ERROR, "Failed to create output context: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    stream_map = av_calloc(pFmtCtx->nb_streams, sizeof(int));
    stream_enabled = av_calloc(pFmtCtx->nb_streams, sizeof(int));
    if(!stream_map || !stream_enabled)
    {
        av_log(NULL, AV_LOG_ERROR, "no memory\n");
        goto _ERROR;
    }

    // 初始化流映射
    for(int i = 0; i < pFmtCtx->nb_streams; i++) {
        stream_map[i] = -1;
        stream_enabled[i] = 0;
    }

    for(int i = 0; i < pFmtCtx->nb_streams; i++)
    {
        AVStream* outStream = NULL;
        AVStream* inStream = pFmtCtx->streams[i];
        AVCodecParameters* inCodecPar = inStream->codecpar;

        if( inCodecPar->codec_type != AVMEDIA_TYPE_AUDIO &&
            inCodecPar->codec_type != AVMEDIA_TYPE_VIDEO)
        {
            continue; // 跳过字幕和其他流，简化处理
        }

        // 4.为目的文件，创建一个新的流
        outStream = avformat_new_stream(oFmtCtx, NULL);
        if(!outStream)
        {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to create output stream\n");
            goto _ERROR;
        }

        ret = avcodec_parameters_copy(outStream->codecpar, inStream->codecpar);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to copy codec parameters: %s\n", av_err2str(ret));
            goto _ERROR;
        }

        // 重要：设置codec_tag为0，让FFmpeg自动选择
        outStream->codecpar->codec_tag = 0;

        // 重要：复制时间基
        outStream->time_base = inStream->time_base;

        // 对于MP4格式，需要确保时间基合理
        if (format_name && strcmp(format_name, "mp4") == 0) {
            if (inCodecPar->codec_type == AVMEDIA_TYPE_VIDEO) {
                // 视频流使用合理的时间基
                if (outStream->time_base.den < 1000) {
                    outStream->time_base = (AVRational){1, 90000};
                }
            } else if (inCodecPar->codec_type == AVMEDIA_TYPE_AUDIO) {
                // 音频流使用采样率作为时间基
                outStream->time_base = (AVRational){1, inCodecPar->sample_rate};
            }
        }

        stream_map[i] = stream_idx++;
        stream_enabled[i] = 1;

        // 记录音视频流索引
        if(inCodecPar->codec_type == AVMEDIA_TYPE_VIDEO && video_stream_index == -1) {
            video_stream_index = i;
        } else if(inCodecPar->codec_type == AVMEDIA_TYPE_AUDIO && audio_stream_index == -1) {
            audio_stream_index = i;
        }

        av_log(NULL, AV_LOG_INFO, "Stream %d mapped to %d, codec: %s, time_base: %d/%d\n",
               i, stream_map[i], avcodec_get_name(inCodecPar->codec_id),
               outStream->time_base.num, outStream->time_base.den);
    }

    // 检查是否有可用的流
    if(stream_idx == 0) {
        av_log(NULL, AV_LOG_ERROR, "No valid audio or video streams found\n");
        goto _ERROR;
    }

    // 绑定
    ret = avio_open(&oFmtCtx->pb, dst, AVIO_FLAG_WRITE);
    if(ret < 0)
    {
        av_log(oFmtCtx, AV_LOG_ERROR, "Failed to open output file: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 7.写多媒体文件头到目标文件中
    AVDictionary* opts = NULL;
    // 对于MP4格式设置优化选项
    if (format_name && strcmp(format_name, "mp4") == 0) {
        av_dict_set(&opts, "movflags", "frag_keyframe+empty_moov", 0);
    }

    ret = avformat_write_header(oFmtCtx, &opts);
    if (opts) av_dict_free(&opts);

    if(ret < 0 )
    {
        av_log(oFmtCtx, AV_LOG_ERROR, "Failed to write header: %s\n", av_err2str(ret));
        goto _ERROR;
    }

    // 改进的定位策略
    int64_t seek_target = (int64_t)(startTime * AV_TIME_BASE);
    int seek_stream = video_stream_index != -1 ? video_stream_index :
                          (audio_stream_index != -1 ? audio_stream_index : 0);

    av_log(NULL, AV_LOG_INFO, "Seeking to %.2f seconds in stream %d\n", startTime, seek_stream);

    // 使用精确搜索
    ret = av_seek_frame(pFmtCtx, seek_stream, seek_target, AVSEEK_FLAG_BACKWARD);
    if(ret < 0)
    {
        av_log(oFmtCtx, AV_LOG_WARNING, "Seek failed: %s, starting from beginning\n", av_err2str(ret));
        // 如果搜索失败，从头开始
        global_start_pts = 0;
    }

    int writing = 0; // 标记是否开始写入
    int packet_count = 0;

    // 8.从源多媒体文件中读取视频/音频/字幕数据到目的文件
    while(av_read_frame(pFmtCtx, &pkt) >= 0)
    {
        int stream_index = pkt.stream_index;

        // 跳过未映射的流
        if(stream_map[stream_index] < 0 || !stream_enabled[stream_index])
        {
            av_packet_unref(&pkt);
            continue;
        }

        AVStream* inStream = pFmtCtx->streams[stream_index];
        AVStream* outStream = oFmtCtx->streams[stream_map[stream_index]];

        // 处理时间戳
        int64_t pts = pkt.pts;
        int64_t dts = pkt.dts;

        // 如果PTS无效，使用DTS
        if (pts == AV_NOPTS_VALUE) {
            pts = dts;
        }
        // 如果都无效，跳过这个包
        if (pts == AV_NOPTS_VALUE) {
            av_packet_unref(&pkt);
            continue;
        }

        // 计算当前时间
        double current_time = pts * av_q2d(inStream->time_base);

        // 检查是否超过结束时间
        if(current_time > endTime)
        {
            av_packet_unref(&pkt);
            break;
        }

        // 检查是否达到开始时间
        if(!writing) {
            if(current_time >= startTime) {
                writing = 1;
                global_start_pts = pts; // 记录第一个有效包的PTS
                av_log(NULL, AV_LOG_INFO, "Start writing at time: %.3f, stream: %d\n", current_time, stream_index);
            } else {
                av_packet_unref(&pkt);
                continue;
            }
        }

        // 调整时间戳（从第一个有效包开始重新计时）
        if(global_start_pts != -1) {
            // 重新计算时间戳（从0开始）
            pkt.pts = pts - global_start_pts;
            pkt.dts = (dts != AV_NOPTS_VALUE) ? dts - global_start_pts : pkt.pts;

            // 确保时间戳非负
            if(pkt.pts < 0) pkt.pts = 0;
            if(pkt.dts < 0) pkt.dts = 0;

            // 确保PTS >= DTS（解决B帧导致的卡帧）
            if(pkt.dts != AV_NOPTS_VALUE && pkt.pts < pkt.dts) {
                pkt.pts = pkt.dts;
            }
        }

        // 更新流索引
        pkt.stream_index = stream_map[stream_index];

        // 时间戳缩放
        av_packet_rescale_ts(&pkt, inStream->time_base, outStream->time_base);

        // 重新计算duration
        if(pkt.duration > 0) {
            pkt.duration = av_rescale_q(pkt.duration, inStream->time_base, outStream->time_base);
        }

        pkt.pos = -1;

        // 写入数据包
        ret = av_interleaved_write_frame(oFmtCtx, &pkt);
        if(ret < 0) {
            av_log(oFmtCtx, AV_LOG_ERROR, "Failed to write frame: %s\n", av_err2str(ret));
        } else {
            packet_count++;
        }

        av_packet_unref(&pkt);

        // 显示进度
        if (packet_count % 100 == 0) {
            av_log(NULL, AV_LOG_INFO, "Progress: %d packets, current time: %.2f\n", packet_count, current_time);
        }
    }

    // 9.写多媒体文件尾到文件中
    av_write_trailer(oFmtCtx);
    av_log(oFmtCtx, AV_LOG_INFO, "Cut successfully completed! Total packets: %d\n", packet_count);

// 10.资源释放
_ERROR:
    if(pFmtCtx)
    {
        avformat_close_input(&pFmtCtx);
    }
    if(oFmtCtx)
    {
        if(oFmtCtx->pb) {
            avio_closep(&oFmtCtx->pb);
        }
        avformat_free_context(oFmtCtx);
    }
    if(stream_map)
    {
        av_free(stream_map);
    }
    if(stream_enabled)
    {
        av_free(stream_enabled);
    }

    return (ret < 0) ? -1 : 0;
}
